"""
AI相关UI扩展
包含AI模式和AI训练的UI功能
"""

import pygame
import os
import subprocess
import threading
from settings import *
from ui import UI  # 导入UI类

class AIUI:
    """AI相关UI类，处理AI模式和AI训练的UI功能"""
    
    def __init__(self, game, main):
        """初始化AI UI"""
        self.game = game
        self.main = main  # Main类的引用
        self.font = pygame.font.SysFont('Arial', 36)
        self.small_font = pygame.font.SysFont('Arial', 24)
        
        # AI对战相关
        self.ai_enabled = False
        self.ai_model_path = None
        self.ai_agent = None
        self.scroll_offset = 0
        
        # 训练相关
        self.training_thread = None
        self.training_process = None
        self.training_active = False
        self.training_progress = 0
        self.training_status = "未开始"
        self.training_output = []
        
        # 模型相关
        self.available_models = []
        self.selected_model = None
        self.update_available_models()
        
    def draw_ai_menu(self):
        """绘制AI模式菜单"""
        self.game.screen.fill(DARKGREY)
        self.game.draw_text("AI 模式", 64, WHITE, WIDTH / 2, 50)
        
        # 绘制AI状态
        ai_status = "启用" if self.ai_enabled else "禁用"
        status_color = GREEN if self.ai_enabled else RED
        self.game.draw_text(f"AI状态: {ai_status}", 36, status_color, WIDTH / 2, 150)
        
        # 如果有加载的模型，显示模型信息
        if self.ai_model_path:
            model_name = os.path.basename(self.ai_model_path)
            self.game.draw_text(f"当前模型: {model_name}", 24, WHITE, WIDTH / 2, 200)
        
        # 绘制按钮
        self.game.toggle_ai_btn = self.game.draw_text("切换AI状态", 36, BLUE, WIDTH / 2, 250)
        self.game.load_model_btn = self.game.draw_text("加载AI模型", 36, YELLOW, WIDTH / 2, 320)
        self.game.ai_training_btn = self.game.draw_text("AI训练", 36, PURPLE, WIDTH / 2, 390)
        self.game.back_btn = self.game.draw_text("返回", 36, RED, WIDTH / 2, 460)
        
        # 绘制操作提示
        self.game.draw_text("在游戏中按A键切换AI状态", 24, LIGHTGREY, WIDTH / 2, 530)
    
    def draw_model_selection_menu(self):
        """绘制AI模型选择菜单"""
        self.game.screen.fill(DARKGREY)
        self.game.draw_text("选择AI模型", 64, WHITE, WIDTH / 2, 50)
        
        # 获取所有模型文件
        model_files = self.get_model_files()
        
        # 计算可见区域
        item_height = 60
        visible_area_height = HEIGHT - 200  # 为标题和按钮留出空间
        items_per_page = max(1, visible_area_height // item_height)
        
        # 计算滚动范围
        total_height = len(model_files) * item_height
        max_scroll = max(0, total_height - visible_area_height)
        
        # 限制滚动偏移量
        self.scroll_offset = max(0, min(self.scroll_offset, max_scroll))
        
        # 显示模型列表
        self.game.model_buttons = []
        start_y = 150 - self.scroll_offset
        
        # 只显示可见的模型文件
        start_index = max(0, self.scroll_offset // item_height)
        end_index = min(len(model_files), start_index + items_per_page + 1)
        
        for i in range(start_index, end_index):
            model_file = model_files[i]
            y_pos = start_y + (i * item_height)
            
            # 只绘制在屏幕范围内的项目
            if y_pos > -item_height and y_pos < HEIGHT:
                # 绘制模型名称
                button = self.game.draw_text(model_file, 36, WHITE, WIDTH / 2, y_pos)
                self.game.model_buttons.append((button, model_file))
        
        # 返回按钮
        self.game.back_btn = self.game.draw_text("返回", 36, RED, WIDTH / 2, HEIGHT - 50)
    
    def draw_training_menu(self):
        """绘制AI训练菜单"""
        self.game.screen.fill(DARKGREY)
        self.game.draw_text("AI 训练", 64, WHITE, WIDTH / 2, 50)
        
        # 训练状态
        if self.training_active:
            status_text = "训练中..."
            status_color = GREEN
        else:
            status_text = "空闲"
            status_color = LIGHTGREY
        
        self.game.draw_text(f"训练状态: {status_text}", 36, status_color, WIDTH / 2, 120)
        
        # 训练参数
        self.game.draw_text("训练参数", 36, WHITE, WIDTH / 2, 180)
        self.game.draw_text(f"算法: PPO", 24, WHITE, WIDTH / 2 - 150, 220)
        self.game.draw_text(f"步数: 100000", 24, WHITE, WIDTH / 2 + 150, 220)
        
        # 绘制按钮
        if not self.training_active:
            self.game.start_training_btn = self.game.draw_text("开始训练", 36, GREEN, WIDTH / 2, 280)
        else:
            self.game.stop_training_btn = self.game.draw_text("停止训练", 36, RED, WIDTH / 2, 280)
        
        self.game.test_model_btn = self.game.draw_text("测试模型", 36, YELLOW, WIDTH / 2, 350)
        self.game.compare_models_btn = self.game.draw_text("比较模型", 36, PURPLE, WIDTH / 2, 420)
        self.game.back_btn = self.game.draw_text("返回", 36, RED, WIDTH / 2, 490)
        
        # 显示训练输出（如果有）
        if self.training_output:
            # 绘制输出框背景
            output_box = pygame.Rect(50, HEIGHT - 150, WIDTH - 100, 100)
            pygame.draw.rect(self.game.screen, (30, 30, 30), output_box)
            pygame.draw.rect(self.game.screen, WHITE, output_box, 1)
            
            # 显示最近的输出
            font = pygame.font.Font(None, 20)
            for i, line in enumerate(self.training_output[-5:]):  # 只显示最近5行
                text_surface = font.render(line, True, WHITE)
                self.game.screen.blit(text_surface, (60, HEIGHT - 140 + i * 20))
    
    def get_model_files(self):
        """获取所有模型文件"""
        model_dir = "models"
        if not os.path.exists(model_dir):
            os.makedirs(model_dir)
            return []
        
        # 获取所有.zip和目录文件
        model_files = []
        for file in os.listdir(model_dir):
            if file.endswith('.zip') or os.path.isdir(os.path.join(model_dir, file)):
                model_files.append(file)
        
        return model_files
    
    def ai_menu_events(self, event):
        """处理AI菜单事件"""
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            # 检查按钮点击
            if hasattr(self.game, 'toggle_ai_btn') and self.game.toggle_ai_btn.collidepoint(event.pos):
                self.toggle_ai()
            elif hasattr(self.game, 'load_model_btn') and self.game.load_model_btn.collidepoint(event.pos):
                self.game.state = 'model_selection'
            elif hasattr(self.game, 'ai_training_btn') and self.game.ai_training_btn.collidepoint(event.pos):
                self.game.state = 'training'
            elif hasattr(self.game, 'back_btn') and self.game.back_btn.collidepoint(event.pos):
                self.game.state = 'main_menu'
    
    def model_selection_events(self, event):
        """处理模型选择菜单事件"""
        if event.type == pygame.MOUSEBUTTONDOWN:
            if event.button == 4:  # 向上滚动
                self.scroll_offset = max(0, self.scroll_offset - 50)
            elif event.button == 5:  # 向下滚动
                # 获取所有模型文件
                model_files = self.get_model_files()
                item_height = 60
                visible_area_height = HEIGHT - 200
                total_height = len(model_files) * item_height
                max_scroll = max(0, total_height - visible_area_height)
                self.scroll_offset = min(max_scroll, self.scroll_offset + 50)
            elif event.button == 1:  # 左键点击
                # 检查按钮点击
                if hasattr(self.game, 'model_buttons'):
                    for btn_rect, model_file in self.game.model_buttons:
                        if btn_rect.collidepoint(event.pos):
                            # 加载选定的模型
                            self.load_model(model_file)
                            self.game.state = 'ai_menu'
                            return
                
                if hasattr(self.game, 'back_btn') and self.game.back_btn.collidepoint(event.pos):
                    self.game.state = 'ai_menu'
    
    def training_menu_events(self, event):
        """处理训练菜单事件"""
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            # 检查按钮点击
            if hasattr(self.game, 'start_training_btn') and self.game.start_training_btn.collidepoint(event.pos):
                if not self.training_active:
                    self.start_training()
            elif hasattr(self.game, 'stop_training_btn') and self.game.stop_training_btn.collidepoint(event.pos):
                if self.training_active:
                    self.stop_training()
            elif hasattr(self.game, 'test_model_btn') and self.game.test_model_btn.collidepoint(event.pos):
                self.test_model()
            elif hasattr(self.game, 'compare_models_btn') and self.game.compare_models_btn.collidepoint(event.pos):
                self.compare_models()
            elif hasattr(self.game, 'back_btn') and self.game.back_btn.collidepoint(event.pos):
                self.game.state = 'ai_menu'
    
    def toggle_ai(self):
        """切换AI状态"""
        self.ai_enabled = not self.ai_enabled
        
        if self.ai_enabled:
            # 启用AI
            if not self.ai_agent:
                from ai_vs_player_mode import AIVsPlayerMode
                self.ai_agent = AIVsPlayerMode(self.main)
            
            # 如果有模型路径，加载模型
            if self.ai_model_path:
                self.ai_agent.set_ai_model(self.ai_model_path)
            else:
                # 使用随机AI
                self.ai_agent.set_ai_model(None)
            
            print("AI已启用")
        else:
            # 禁用AI
            self.ai_agent = None
            print("AI已禁用")
    
    def load_model(self, model_file):
        """加载AI模型"""
        model_path = os.path.join("models", model_file)
        if os.path.exists(model_path):
            self.ai_model_path = model_path
            print(f"已加载AI模型: {model_path}")
            
            # 如果AI已启用，更新AI模型
            if self.ai_enabled and self.ai_agent:
                self.ai_agent.set_ai_model(model_path)
        else:
            print(f"模型文件不存在: {model_path}")
    
    def start_training(self):
        """开始训练"""
        if self.training_active:
            return
        
        self.training_active = True
        self.training_output = ["开始训练..."]
        
        # 在新线程中运行训练，避免阻塞UI
        def run_training():
            try:
                # 运行训练脚本
                process = subprocess.Popen(
                    ["python", "train_ai.py", "--mode", "train", "--algorithm", "PPO", "--timesteps", "100000"],
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    text=True,
                    bufsize=1,
                    universal_newlines=True
                )
                
                self.training_process = process
                
                # 读取输出
                for line in iter(process.stdout.readline, ''):
                    if line:
                        self.training_output.append(line.strip())
                        # 只保留最近的20行输出
                        if len(self.training_output) > 20:
                            self.training_output.pop(0)
                
                # 等待进程完成
                process.wait()
                self.training_output.append(f"训练完成，退出码: {process.returncode}")
            except Exception as e:
                self.training_output.append(f"训练出错: {str(e)}")
            finally:
                self.training_active = False
                self.training_process = None
        
        # 启动训练线程
        training_thread = threading.Thread(target=run_training)
        training_thread.daemon = True
        training_thread.start()
    
    def stop_training(self):
        """停止训练"""
        if self.training_process:
            self.training_process.terminate()
            self.training_output.append("训练已停止")
        self.training_active = False
    
    def test_model(self):
        """测试模型"""
        # 获取最新的模型文件
        model_files = self.get_model_files()
        if not model_files:
            self.training_output.append("没有可用的模型文件")
            return
        
        # 使用最新的模型文件
        latest_model = max(model_files, key=lambda f: os.path.getmtime(os.path.join("models", f)))
        model_path = os.path.join("models", latest_model)
        
        # 在新线程中运行测试，避免阻塞UI
        def run_test():
            try:
                self.training_output.append(f"开始测试模型: {latest_model}")
                
                # 运行测试脚本
                process = subprocess.Popen(
                    ["python", "train_ai.py", "--mode", "test", "--model", model_path, "--episodes", "3"],
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    text=True,
                    bufsize=1,
                    universal_newlines=True
                )
                
                # 读取输出
                for line in iter(process.stdout.readline, ''):
                    if line:
                        self.training_output.append(line.strip())
                        # 只保留最近的20行输出
                        if len(self.training_output) > 20:
                            self.training_output.pop(0)
                
                # 等待进程完成
                process.wait()
                self.training_output.append(f"测试完成，退出码: {process.returncode}")
            except Exception as e:
                self.training_output.append(f"测试出错: {str(e)}")
        
        # 启动测试线程
        test_thread = threading.Thread(target=run_test)
        test_thread.daemon = True
        test_thread.start()
    
    def compare_models(self):
        """比较模型"""
        # 获取最新的模型文件
        model_files = self.get_model_files()
        if not model_files:
            self.training_output.append("没有可用的模型文件")
            return
        
        # 使用最新的模型文件
        latest_model = max(model_files, key=lambda f: os.path.getmtime(os.path.join("models", f)))
        model_path = os.path.join("models", latest_model)
        
        # 在新线程中运行比较，避免阻塞UI
        def run_compare():
            try:
                self.training_output.append(f"开始比较模型: {latest_model}")
                
                # 运行比较脚本
                process = subprocess.Popen(
                    ["python", "train_ai.py", "--mode", "compare", "--model", model_path, "--episodes", "5"],
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    text=True,
                    bufsize=1,
                    universal_newlines=True
                )
                
                # 读取输出
                for line in iter(process.stdout.readline, ''):
                    if line:
                        self.training_output.append(line.strip())
                        # 只保留最近的20行输出
                        if len(self.training_output) > 20:
                            self.training_output.pop(0)
                
                # 等待进程完成
                process.wait()
                self.training_output.append(f"比较完成，退出码: {process.returncode}")
            except Exception as e:
                self.training_output.append(f"比较出错: {str(e)}")
        
        # 启动比较线程
        compare_thread = threading.Thread(target=run_compare)
        compare_thread.daemon = True
        compare_thread.start()
    
    def update_in_game(self, dt):
        """在游戏中更新AI状态"""
        if self.ai_enabled and self.ai_agent:
            self.ai_agent.update(dt)
    
    def draw_in_game_ui(self):
        """在游戏中绘制AI状态UI"""
        if not self.ai_enabled:
            return
        
        # 绘制AI状态文本
        font = pygame.font.Font(None, 24)
        ai_status = "AI: 启用"
        if self.ai_model_path:
            model_name = os.path.basename(self.ai_model_path)
            ai_status += f" ({model_name})"
        
        text_surface = font.render(ai_status, True, GREEN)
        self.game.screen.blit(text_surface, (10, 10))
        
        # 绘制控制提示
        help_text = "按A键切换AI状态"
        help_surface = font.render(help_text, True, WHITE)
        self.game.screen.blit(help_surface, (10, 40))
    
    def update_available_models(self):
        """更新可用模型列表"""
        self.available_models = self.get_model_files()
    
    def run_ai_menu(self):
        """运行AI菜单"""
        self.game.playing = True
        while self.game.playing and self.game.state == 'ai_menu':
            self.game.dt = self.game.clock.tick(self.game.FPS) / 1000.0
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    exit()
                self.ai_menu_events(event)
            
            self.draw_ai_menu()
            pygame.display.flip()
    
    def run_model_selection(self):
        """运行模型选择菜单"""
        self.game.playing = True
        while self.game.playing and self.game.state == 'model_selection':
            self.game.dt = self.game.clock.tick(self.game.FPS) / 1000.0
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    exit()
                self.model_selection_events(event)
            
            self.draw_model_selection_menu()
            pygame.display.flip()
    
    def run_training(self):
        """运行训练菜单"""
        self.game.playing = True
        while self.game.playing and self.game.state == 'training':
            self.game.dt = self.game.clock.tick(self.game.FPS) / 1000.0
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    exit()
                self.training_events(event)
            
            self.draw_training_menu()
            pygame.display.flip()
    
    def run_ai_vs_player(self):
        """运行AI对战模式"""
        while self.game.state == 'ai_vs_player':
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    exit()
                self.ai_vs_player_events(event)
            
            self.draw_ai_vs_player()
            pygame.display.flip()
            self.game.clock.tick(60)
    
    def training_events(self, event):
        """处理训练菜单事件"""
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            # 检查按钮点击
            if hasattr(self.game, 'start_training_btn') and self.game.start_training_btn.collidepoint(event.pos):
                if not self.training_active:
                    self.start_training()
            elif hasattr(self.game, 'stop_training_btn') and self.game.stop_training_btn.collidepoint(event.pos):
                if self.training_active:
                    self.stop_training()
            elif hasattr(self.game, 'test_model_btn') and self.game.test_model_btn.collidepoint(event.pos):
                self.test_model()
            elif hasattr(self.game, 'compare_models_btn') and self.game.compare_models_btn.collidepoint(event.pos):
                self.compare_models()
            elif hasattr(self.game, 'back_btn') and self.game.back_btn.collidepoint(event.pos):
                self.game.state = 'ai_menu'
    
    def ai_vs_player_events(self, event):
        """处理AI对战事件"""
        # 这里应该处理AI对战模式的事件
        # 暂时简化处理，直接返回主菜单
        if event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE:
            self.game.state = 'main_menu'
        elif event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            # 检查返回按钮点击
            # 使用draw_ai_vs_player返回的按钮矩形
            back_btn = self.draw_ai_vs_player()
            if back_btn and back_btn.collidepoint(event.pos):
                self.game.state = 'main_menu'
    
    def draw_ai_vs_player(self):
        """绘制AI对战界面"""
        self.game.screen.fill(DARKGREY)
        self.game.draw_text("AI 对战模式", 64, WHITE, WIDTH / 2, 50)
        
        # 显示提示信息
        self.game.draw_text("按ESC返回主菜单", 36, WHITE, WIDTH / 2, HEIGHT / 2)
        
        # 返回按钮
        back_btn = self.game.draw_text("返回", 36, WHITE, WIDTH / 2, HEIGHT - 50)
        return back_btn